#!/usr/bin/env python3
"""Rename common English column names in a Parquet file to Chinese equivalents.

Usage:
    python scripts/rename_parquet_columns.py /path/to/file.parquet

The script creates a timestamped backup under the same folder `backup/` before writing.
Unknown columns are left unchanged and listed in the log.
"""
from __future__ import annotations

import sys
from pathlib import Path
from datetime import datetime
import shutil

try:
    import polars as pl
except Exception:
    pl = None

try:
    import pandas as pd
except Exception:
    pd = None


DEFAULT_MAPPING = {
    "Year": "年份",
    "Prvn": "省",
    "Pftn": "地级市",
    "Cont": "县",
    "Vlgn": "乡镇",
}


def backup(src: Path) -> Path:
    backup_dir = src.parent / "backup"
    backup_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%dT%H%M%S")
    dest = backup_dir / f"{src.name}.{ts}.bak"
    shutil.copy2(src, dest)
    return dest


def rename_with_polars(path: Path, mapping: dict) -> int:
    df = pl.read_parquet(path)
    cols_before = df.columns
    # Build rename map for existing columns only
    rename_map = {c: mapping[c] for c in cols_before if c in mapping}
    if rename_map:
        df = df.rename(rename_map)
    df.write_parquet(path)
    return df.height


def rename_with_pandas(path: Path, mapping: dict) -> int:
    df = pd.read_parquet(path)
    cols_before = list(df.columns)
    rename_map = {c: mapping[c] for c in cols_before if c in mapping}
    if rename_map:
        df = df.rename(columns=rename_map)
    df.to_parquet(path, index=False)
    return df.shape[0]


def main():
    if len(sys.argv) < 2:
        print("Usage: rename_parquet_columns.py /path/to/file.parquet")
        sys.exit(2)

    target = Path(sys.argv[1]).resolve()
    if not target.exists():
        print(f"File not found: {target}")
        sys.exit(2)

    print(f"Target file: {target}")
    # read columns only using polars if available
    try:
        if pl is not None:
            cols = pl.read_parquet(target, n_rows=0).columns
        elif pd is not None:
            cols = list(pd.read_parquet(target, engine='pyarrow').columns)
        else:
            print("Neither polars nor pandas available to inspect parquet file.")
            sys.exit(1)
    except Exception as e:
        print(f"Failed to inspect file columns: {e}")
        sys.exit(1)

    print("Columns before:")
    for c in cols:
        print(" -", c)

    mapping = DEFAULT_MAPPING
    applied = {c: mapping[c] for c in cols if c in mapping}
    unmapped = [c for c in cols if c not in mapping]

    print("Planned renames:")
    if applied:
        for k, v in applied.items():
            print(f" - {k} -> {v}")
    else:
        print(" - (no matching columns to rename)")

    if unmapped:
        print("Columns left unchanged:")
        for c in unmapped:
            print(" -", c)

    # Backup original
    try:
        backup_path = backup(target)
        print(f"Backed up original to: {backup_path}")
    except Exception as e:
        print(f"Backup failed: {e}")
        sys.exit(1)

    # Perform rename and write
    try:
        if pl is not None:
            new_count = rename_with_polars(target, mapping)
            print(f"Wrote renamed parquet with {new_count} rows (Polars).")
        elif pd is not None:
            new_count = rename_with_pandas(target, mapping)
            print(f"Wrote renamed parquet with {new_count} rows (pandas).")
        else:
            print("No engine available to write parquet.")
            shutil.copy2(backup_path, target)
            sys.exit(1)
    except Exception as e:
        print(f"Rename/write failed: {e}")
        # try to restore
        try:
            shutil.copy2(backup_path, target)
            print("Original restored from backup.")
        except Exception as e2:
            print(f"Failed to restore backup: {e2}")
        sys.exit(1)


if __name__ == "__main__":
    main()
